/**
 * 
 */
package com.zxc.redis;

import java.util.List;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

/**
 * @author Administrator
 *
 */
public class TokenBucket {

	private final JedisPool jedisPool;
    private final String script;

    public TokenBucket(JedisPool jedisPool) {
        this.jedisPool = jedisPool;
        // Lua脚本内容
        this.script = "local key = KEYS[1] " +
                      "local rate = tonumber(ARGV[1]) " +
                      "local capacity = tonumber(ARGV[2]) " +
                      "local now = tonumber(ARGV[3]) " +
                      "local requested = tonumber(ARGV[4]) " +
                      "local last_tokens = tonumber(redis.call('get', key .. ':tokens')) or capacity " +
                      "local last_refreshed = tonumber(redis.call('get', key .. ':timestamp')) or now " +
                      "local delta = math.max(0, now - last_refreshed) * rate / 1000 " +
                      "local filled_tokens = math.min(capacity, last_tokens + delta) " +
                      "if filled_tokens < requested then " +
                      "    return {0, filled_tokens} " +
                      "else " +
                      "    filled_tokens = filled_tokens - requested " +
                      "    redis.call('set', key .. ':tokens', filled_tokens) " +
                      "    redis.call('set', key .. ':timestamp', now) " +
                      "    return {1, filled_tokens} " +
                      "end";
    }

    @SuppressWarnings("unchecked")
	public boolean acquireToken(String key, int rate, int capacity, int requested) {
        try (Jedis jedis = jedisPool.getResource()) {
            long now = System.currentTimeMillis();
            Object result = jedis.eval(script, 1, key, String.valueOf(rate), String.valueOf(capacity), String.valueOf(now), String.valueOf(requested));
            List<Long> results = (List<Long>) result;
            return results.get(0) == 1;
        }
    }
}
